#include "torrent.h"

/*******************************************************************************
 * Constructors
 *******************************************************************************/
Torrent::Torrent(const char *filename)
{
	// open file
	int fd;
	if ((fd = open(filename, O_RDONLY)) == -1)
	{
		printf("Failed to open torrent file.\n");
		exit(-1);
		return;
	}

	// determine file size
	struct stat stats;
	fstat(fd, &stats);
	int filesize = stats.st_size;

	// read whole torrent file into memory
	char *buffer = (char*)malloc(sizeof(char)*filesize);
	read(fd, buffer, filesize);
	close(fd);

	// bdecode the buffer
	bDictionary *dictionary = (bDictionary*)bDecode(buffer);

	// get the announce info
	string announce = ((bString*)dictionary->GetItemFromKey("announce"))->GetString();
	if (announce.find("http://") != string::npos)
	{
		string tracker_hostname = announce.substr(7, announce.find("/", 7) - 7);
		if (tracker_hostname.find_last_of(":") > 7)
		{
			tracker_hostname = tracker_hostname.substr(0, tracker_hostname.find_last_of(":"));
		}
		string announce_dir = announce.substr(announce.find("/", 7), announce.length()-announce.find("/"));
		unsigned tracker_port = 0;
		if (announce.find(":", 7) != string::npos)
		{
			string tmp = announce.substr(announce.find(":", 7)+1, announce.find("/", 7)-announce.find(":", 7));
			tracker_port = atoi(tmp.c_str());
		}
		else
		{
			tracker_port = 80;
		}
		m_trackers.push_back(new Tracker(this, tracker_hostname, tracker_port, announce_dir));
	}
	bList *announce_list = (bList*)dictionary->GetItemFromKey("announce-list");
	if (announce_list != NULL)
	{
		for (unsigned i = 0; i < announce_list->GetNumItems(); i++)
		{
			bList *announce_group = (bList*)announce_list->GetItemFromIndex(i);
			for (unsigned j = 0; j < announce_group->GetNumItems(); j++)
			{
				announce = ((bString*)announce_group->GetItemFromIndex(j))->GetString();
				if (announce.find("http://") != string::npos)
				{
					string tracker_hostname = announce.substr(7, announce.find("/", 7) - 7);
					if (tracker_hostname.find_last_of(":") > 7)
					{
						tracker_hostname = tracker_hostname.substr(0, tracker_hostname.find_last_of(":"));
					}
					string announce_dir = announce.substr(announce.find("/", 7), announce.length()-announce.find("/"));
					unsigned tracker_port = 0;
					if (announce.find(":", 7) != string::npos)
					{
						string tmp = announce.substr(announce.find(":", 7)+1, announce.find("/", 7)-announce.find(":", 7));
						tracker_port = atoi(tmp.c_str());
					}
					else
					{
						tracker_port = 80;
					}
					m_trackers.push_back(new Tracker(this, tracker_hostname, tracker_port, announce_dir));
				}
			}
		}
	}
	if (m_trackers.size() == 0)
	{
		printf("Torrents with only non-HTTP trackers are not supported.\n");
		free(buffer);
		exit(-1);
		return;
	}

	// get other info
	bDictionary *info = (bDictionary*)dictionary->GetItemFromKey("info");
	if (info->GetItemFromKey("length") == NULL)
	{
		printf("Multi-file torrents are not supported.\n");
		free(buffer);
		exit(-1);
		return;
	}
	bInteger *prvt = (bInteger*)info->GetItemFromKey("private");
	if (prvt != NULL && (bool)(prvt)->GetInteger())
	{
		printf("Private trackers are not supported.\n");
		free(buffer);
		exit(-1);
		return;
	}
	m_filename = ((bString*)info->GetItemFromKey("name"))->GetString();
	m_filesize = ((bInteger*)info->GetItemFromKey("length"))->GetInteger();
	m_piecesize = ((bInteger*)info->GetItemFromKey("piece length"))->GetInteger();
	m_numpieces = ceilf(((float)m_filesize)/((float)m_piecesize));

	// get piece info
	m_pieces = (unsigned char*)malloc(sizeof(unsigned char)*20*m_numpieces);
	((bString*)((bDictionary*)dictionary->GetItemFromKey("info"))->GetItemFromKey("pieces"))->GetByteString((char*)m_pieces);

	// set up the torrent struct
	int offset = -1;
	while (strncmp(&(buffer[++offset]), "info", 4));
	offset += 4;

	// compute sha1 hash for this torrent file
	SHA1Context sha;
	SHA1Reset(&sha);
	SHA1Input(&sha, (const unsigned char *) &(buffer[offset]), filesize-1-offset);
	if (!SHA1Result(&sha))
	{
		printf("Failed to compute message digest\n");
		exit(-1);
	}
	else
	{
		for(int i = 0; i < 5; i++)
		{
			m_infohash[i*4+0] = ((char*)&(sha.Message_Digest[i]))[3];
			m_infohash[i*4+1] = ((char*)&(sha.Message_Digest[i]))[2];
			m_infohash[i*4+2] = ((char*)&(sha.Message_Digest[i]))[1];
			m_infohash[i*4+3] = ((char*)&(sha.Message_Digest[i]))[0];
		}
	}

	// initialize download statistics
	m_uploaded = 0;
	m_downloaded = 0;

	// initialize bitfield
	unsigned fieldsize = ceilf((float)(m_numpieces)/8.0f);
	m_bitfield = (unsigned char*)malloc(sizeof(unsigned char)*fieldsize);
	memset(m_bitfield, 0, fieldsize);

	//create the storage to save the pieces to
	m_storage = new Storage(m_filename, m_piecesize, m_numpieces, m_filesize);

	//save bitfield based on what we already have
	DoResume();

	// free memory
	delete dictionary;
	free(buffer);
}


void Torrent::DoResume()
{
	unsigned pieces = GetNumPieces();
	unsigned filelen = GetFileSize();
	unsigned piecelen = GetPieceSize();
	int count = 0;
	for(unsigned i=0;i<pieces;i++)
	{
		unsigned thispsize = piecelen;
		if(i==pieces-1)
			thispsize = (unsigned) (filelen  - (pieces-1)*piecelen);
		Piece p(i,thispsize);
		unsigned char *data = (unsigned char*)malloc(sizeof(unsigned char)*thispsize+1);
		m_storage->GetPiece(data,i,thispsize);
		p.SetData(data,thispsize);
		if(p.CompareHashValue(GetPieceHash(i)))
		{
			count++;
			//			cout<<"we have piece "<<i<<endl;
			SetHavePiece(i);
		}
		else
		{
			free(data);
			break;
		}
		free(data);
	}
	if (count > 0)
		cout << "Resuming with " << count << "/" << pieces << " pieces complete." << endl;
}

Torrent::~Torrent()
{
//	cout<<"torrent deleted"<<endl;
	for(int i = 0;i<(int)m_peers.size();i++)
	{
		delete m_peers.at(i);
	}
	for(int i = 0;i<(int)m_trackers.size();i++)
	{
		delete m_trackers.at(i);
	}
	if (m_pieces != NULL)
		free(m_pieces);
	if (m_bitfield != NULL)
		free(m_bitfield);
	if (m_storage != NULL)
		delete m_storage;
}

/*******************************************************************************
 * Torrent Data
 *******************************************************************************/
string Torrent::GetFileName()
{
	return m_filename;
}
unsigned Torrent::GetFileSize()
{
	return m_filesize;
}
unsigned Torrent::GetPieceSize(unsigned piece_num)
{
	unsigned ret;
	if (piece_num < m_numpieces-1)
		ret = m_piecesize;
	else
		ret = m_filesize - m_piecesize*(m_numpieces-1);
	return ret;
}
unsigned Torrent::GetNumPieces()
{
	return m_numpieces;
}
unsigned Torrent::GetUploaded()
{
	return m_uploaded * m_piecesize;
}
unsigned Torrent::GetDownloaded()
{
	if (m_downloaded == m_numpieces)
		return m_filesize;
	return m_downloaded * m_piecesize;
}

string Torrent::GetInfoHash()
{
	char buffer[41];
	for(int i = 0; i < 5 ; i++)
	{
		unsigned flip;
		((char*)&flip)[0] = m_infohash[i*4+3];
		((char*)&flip)[1] = m_infohash[i*4+2];
		((char*)&flip)[2] = m_infohash[i*4+1];
		((char*)&flip)[3] = m_infohash[i*4+0];
		sprintf(&(buffer[i*8]), "%X", flip);
	}
	buffer[40] = '\0';
	return string(buffer);
}

/*******************************************************************************
 * Peer/Tracker Data
 *******************************************************************************/
unsigned Torrent::GetNumPeers()
{
	return m_peers.size();
}
unsigned Torrent::GetNumTrackers()
{
	return m_trackers.size();
}
Peer* Torrent::GetPeer(unsigned index)
{
	return m_peers[index];
}
Tracker* Torrent::GetTracker(unsigned index)
{
	return m_trackers[index];
}

/*******************************************************************************
 * Peer/Tracker use only
 *******************************************************************************/
unsigned char* Torrent::GetRawInfoHash()
{
	return m_infohash;
}
unsigned char* Torrent::GetPieceHash(unsigned index)
{
	return &(m_pieces[index*20]);
}
Storage* Torrent::GetStorage()
{
	return m_storage;
}
void Torrent::AddPeer(Peer *peer)
{
	m_peers.push_back(peer);
}
unsigned Torrent::GetIndexOfNextPiece()
{
	unsigned ret;
	for (ret = 0; ret < m_numpieces; ret++)
	{
		unsigned byte = ret / 8;
		unsigned bit = ret % 8;
		if (m_bitfield[byte] & (1<<(8-bit-1)))
			continue;
		break;
	}
	return ret;
}

bool Torrent::SetHavePiece(unsigned index)
{
	unsigned byte = index/8;
	unsigned bit = index%8;
	unsigned char *spot = m_bitfield + byte;
	char mask = 1;
	mask = mask << (8-bit-1);

	if (!(*spot & mask))
		m_downloaded += 1;
	*spot = *spot | mask;
	return true;
}
